/**
 * Created by Administrator on 2025/5/15.
 * */
#include "AIE2ENet.h"
#include <cuda_runtime.h>
#include <iostream>
#include <thread>
#include "../include/zy_net.h"

BaseObject *createAIE2ENet() {
    return new AIE2ENet();
}

void destroyAIE2ENet(BaseObject *p) {
    SAFE_DELETE(p);
}

/*__global__ void TaskKernel(char *in, char *out) {
    int tid = blockIdx.x * blockDim.x + threadIdx.x;
    out[tid] = in[tid];
}*/

AIE2ENet::AIE2ENet(){
    this->setGPUDevice(0);
    prevStep = -1;
    // 订阅数据源主题Topic
    sub = zmq::socket_t(ctx, zmq::socket_type::sub);
    sub.connect(topicAddr.c_str());
    sub.set(zmq::sockopt::subscribe, topicName);
    // AIE2E通知通道
    req_AIE2E_socket = zmq::socket_t(ctx, zmq::socket_type::req);
    req_AIE2E_socket.connect(req_AIE2E_address);
    // 接收ZMQ请求
    rep_socket = zmq::socket_t(ctx, zmq::socket_type::rep);
    rep_socket.bind(rep_address);
    std::thread thd(&AIE2ENet::zmqRepServe, this);
    thd.detach();
}

AIE2ENet::~AIE2ENet() {
}

void AIE2ENet::zmqRepServe() {
    // 处理CFAR控制命令
    zmq::message_t command;
    try {
        rep_socket.recv(command, zmq::recv_flags::none);
    } catch (const zmq::error_t& e) {
        std::cerr << "ZMQ error: " << e.what() << std::endl;
    }
    int state;
    memcpy(&state, command.data(), sizeof(int));
    if (1 == state) {
        std::cout << "AIE2ETrain::AIE2ENet：开始模型训练 ^_^" << std::endl;
    }
    // 发送确认
    int ackVal = 200;
    zmq::message_t ack(sizeof(int));
    memcpy(ack.data(), &ackVal, sizeof(int));
    rep_socket.send(ack, zmq::send_flags::none);
}

void AIE2ENet::onIssueStreamedCudaOperations(buffer_table_t *input, buffer_table_t *output, cudaStream_t stream) {
    int step;
    zmq::message_t msgTopicName;
    zmq::message_t msgData;
    if (!sub.recv(msgTopicName, zmq::recv_flags::none)) {
        std::cerr << "Subscribe: Failed to receive topic name" << std::endl;
    }
    std::string topicName_(static_cast<char*>(msgTopicName.data()), msgTopicName.size());
    if (!sub.recv(msgData, zmq::recv_flags::none)) {
        std::cerr << "Subscribe: Failed to receive data" << std::endl;
    }
    memcpy(&step, msgData.data(), sizeof(float));
    if (prevStep == step) {
        return ;
    }
    prevStep = step;
    int aiE2ENetState = GlobalParameters::shareInstance()->get_parameter_aiE2ENetState(dynamic_cast<TaskModule*>(this));
    // 初始训练过程
    if (aiE2ENetState==static_cast<int>(AIE2ENetStateEnum::Notify_Initial)) {
        std::cout << "AIE2ETrain.AIE2ENet: 开始训练模型..." << std::endl;
        auto aiE2ENetStateHolder = GlobalParameters::shareInstance()->findParameter("aiE2ENetState");
        int *aiE2ENetState = static_cast<int*>(GlobalParameters::shareInstance()->getValue(aiE2ENetStateHolder, dynamic_cast<TaskModule*>(this)));
        *aiE2ENetState = static_cast<int>(AIE2ENetStateEnum::Initial_Training);
    }
    if (aiE2ENetState==static_cast<int>(AIE2ENetStateEnum::Initial_Training) && step==STEP3_2) {
        std::cout << "AIE2ETrain.AIE2ENet: 模型训练完成，通知AIE2EInfer::AIE2E模块升级模型" << std::endl;
        auto aiE2ENetStateHolder = GlobalParameters::shareInstance()->findParameter("aiE2ENetState");
        int *aiE2ENetState = static_cast<int*>(GlobalParameters::shareInstance()->getValue(aiE2ENetStateHolder, dynamic_cast<TaskModule*>(this)));
        *aiE2ENetState = static_cast<int>(AIE2ENetStateEnum::Initial_Trained);
        // 通知AIE2EInfer::AIE2EInferController更新全局状态aiE2EState，从而使AIE2EInfer::AIE2E模块升级模型
        int aiE2EStateVal = 1;
        cudaMemcpyAsync(output->list[1]->data, &aiE2EStateVal, sizeof(int), cudaMemcpyHostToDevice, stream);
    }
    if (aiE2ENetState==static_cast<int>(AIE2ENetStateEnum::Initial_Trained) && step>=STEP3_2) {
        // 通知AIE2EInfer::AIE2EInferController更新全局状态aiE2EState，从而使AIE2EInfer::AIE2E模块升级模型
        int aiE2EStateVal = 0;
        cudaMemcpyAsync(output->list[1]->data, &aiE2EStateVal, sizeof(int), cudaMemcpyHostToDevice, stream);
    }
    // 传统算法协助训练
    if (aiE2ENetState==static_cast<int>(AIE2ENetStateEnum::Tat_Notify)) {
        std::cout << "AIE2ETrain::AIE2ENet: 开始训练模型（传统算法辅助）..." << std::endl;
        auto aiE2ENetStateHolder = GlobalParameters::shareInstance()->findParameter("aiE2ENetState");
        int *aiE2ENetState = static_cast<int*>(GlobalParameters::shareInstance()->getValue(aiE2ENetStateHolder, dynamic_cast<TaskModule*>(this)));
        *aiE2ENetState = static_cast<int>(AIE2ENetStateEnum::Tat_Training);        
    }
    if (aiE2ENetState==static_cast<int>(AIE2ENetStateEnum::Tat_Training) && step==STEP8) {
        std::cout << "AIE2ETrain::AIE2ENet: 模型训练完成（传统算法辅助），通知AIE2EInfer::AIE2E模块升级模型" << std::endl;
        auto aiE2ENetStateHolder = GlobalParameters::shareInstance()->findParameter("aiE2ENetState");
        int *aiE2ENetState = static_cast<int*>(GlobalParameters::shareInstance()->getValue(aiE2ENetStateHolder, dynamic_cast<TaskModule*>(this)));
        *aiE2ENetState = static_cast<int>(AIE2ENetStateEnum::Tat_Trained);
        // 通知AIE2EInfer::AIE2EInferController更新全局状态aiE2EState，从而使AIE2EInfer::AIE2E模块升级模型
        int aiE2EStateVal = 4;
        cudaMemcpyAsync(output->list[1]->data, &aiE2EStateVal, sizeof(int), cudaMemcpyHostToDevice, stream);
    }
    if (aiE2ENetState==static_cast<int>(AIE2ENetStateEnum::Tat_Trained) && step>STEP8) {
        // 通知AIE2EInfer::AIE2EInferController更新全局状态aiE2EState，从而使AIE2EInfer::AIE2E模块升级模型
        int aiE2EStateVal = 0;
        cudaMemcpyAsync(output->list[1]->data, &aiE2EStateVal, sizeof(int), cudaMemcpyHostToDevice, stream);
    }
    // 虚拟真值协助训练
    if (aiE2ENetState==static_cast<int>(AIE2ENetStateEnum::Vtt_Notify)) {
        std::cout << "AIE2ETrain::AIE2ENet: 开始训练模型（虚拟真值）..." << std::endl;
        auto aiE2ENetStateHolder = GlobalParameters::shareInstance()->findParameter("aiE2ENetState");
        int *aiE2ENetState = static_cast<int*>(GlobalParameters::shareInstance()->getValue(aiE2ENetStateHolder, dynamic_cast<TaskModule*>(this)));
        *aiE2ENetState = static_cast<int>(AIE2ENetStateEnum::Vtt_Training);        
    }
    if (aiE2ENetState==static_cast<int>(AIE2ENetStateEnum::Vtt_Training) && step==STEP11) {
        std::cout << "AIE2ETrain::AIE2ENet: 模型训练完成（虚拟真值），通知AIE2EInfer::AIE2E模块升级模型" << std::endl;
        auto aiE2ENetStateHolder = GlobalParameters::shareInstance()->findParameter("aiE2ENetState");
        int *aiE2ENetState = static_cast<int*>(GlobalParameters::shareInstance()->getValue(aiE2ENetStateHolder, dynamic_cast<TaskModule*>(this)));
        *aiE2ENetState = static_cast<int>(AIE2ENetStateEnum::Vtt_Trained);
        // 通知AIE2EInfer::AIE2EInferController更新全局状态aiE2EState，从而使AIE2EInfer::AIE2E模块升级模型
        int aiE2EStateVal = 7;
        cudaMemcpyAsync(output->list[1]->data, &aiE2EStateVal, sizeof(int), cudaMemcpyHostToDevice, stream);
    }
    if (aiE2ENetState==static_cast<int>(AIE2ENetStateEnum::Vtt_Trained) && step>STEP11) {
        // 通知AIE2EInfer::AIE2EInferController更新全局状态aiE2EState，从而使AIE2EInfer::AIE2E模块升级模型
        int aiE2EStateVal = 0;
        cudaMemcpyAsync(output->list[1]->data, &aiE2EStateVal, sizeof(int), cudaMemcpyHostToDevice, stream);
    }
    // 自监督学习
    if (aiE2ENetState==static_cast<int>(AIE2ENetStateEnum::Sst_Notify)) {
        std::cout << "AIE2ETrain::AIE2ENet: 开始训练模型（自监督学习）..." << std::endl;
        auto aiE2ENetStateHolder = GlobalParameters::shareInstance()->findParameter("aiE2ENetState");
        int *aiE2ENetState = static_cast<int*>(GlobalParameters::shareInstance()->getValue(aiE2ENetStateHolder, dynamic_cast<TaskModule*>(this)));
        *aiE2ENetState = static_cast<int>(AIE2ENetStateEnum::Sst_Training);        
    }
    if (aiE2ENetState==static_cast<int>(AIE2ENetStateEnum::Sst_Training) && step==STEP14) {
        std::cout << "AIE2ETrain::AIE2ENet: 模型训练完成（自监督学习），通知AIE2EInfer::AIE2E模块升级模型" << std::endl;
        auto aiE2ENetStateHolder = GlobalParameters::shareInstance()->findParameter("aiE2ENetState");
        int *aiE2ENetState = static_cast<int*>(GlobalParameters::shareInstance()->getValue(aiE2ENetStateHolder, dynamic_cast<TaskModule*>(this)));
        *aiE2ENetState = static_cast<int>(AIE2ENetStateEnum::Sst_Trained);
        // 通知AIE2EInfer::AIE2EInferController更新全局状态aiE2EState，从而使AIE2EInfer::AIE2E模块升级模型
        int aiE2EStateVal = 10;
        cudaMemcpyAsync(output->list[1]->data, &aiE2EStateVal, sizeof(int), cudaMemcpyHostToDevice, stream);
    }
    if (aiE2ENetState==static_cast<int>(AIE2ENetStateEnum::Sst_Trained) && step>STEP14) {
        // 通知AIE2EInfer::AIE2EInferController更新全局状态aiE2EState，从而使AIE2EInfer::AIE2E模块升级模型
        int aiE2EStateVal = 0;
        cudaMemcpyAsync(output->list[1]->data, &aiE2EStateVal, sizeof(int), cudaMemcpyHostToDevice, stream);
    }
}
